from pymilvus import connections, Collection
from sentence_transformers import SentenceTransformer

# --- 配置参数 ---

# Milvus 连接配置
MILVUS_URI = "http://111.229.125.24:19530"
MILVUS_DB_NAME = "customer_service"
MILVUS_TOKEN = "root:Milvus"
COLLECTION_NAME = "general_faq"

# 用于生成向量的预训练模型 (必须与导入时使用的模型一致)
MODEL_NAME = 'paraphrase-multilingual-MiniLM-L12-v2'

# 搜索参数
# 'nprobe' 是IVF系列索引的查询参数，值越大，召回率越高，但速度越慢。
# 需要根据您的数据量和对精度的要求进行调整。
SEARCH_PARAMS = {
    "metric_type": "L2",  # L2是欧氏距离，值越小表示越相似
    "params": {"nprobe": 64},
}
TOP_K = 1 # 设置为您希望返回的最相似结果的数量

# --- 主逻辑 ---

def main():
    """
    主函数：连接Milvus，加载模型，并循环等待用户输入进行搜索。
    """
    # 1. 连接到Milvus
    print(f"正在连接到 Milvus (URI: {MILVUS_URI})...")
    try:
        connections.connect(
            uri=MILVUS_URI,
            db_name=MILVUS_DB_NAME,
            token=MILVUS_TOKEN
        )
        print("成功连接到 Milvus。")
        collection = Collection(COLLECTION_NAME)
    except Exception as e:
        print(f"连接 Milvus 时发生错误: {e}")
        return

    # 2. 加载向量化模型
    print(f"正在加载向量模型 '{MODEL_NAME}'...")
    try:
        model = SentenceTransformer(MODEL_NAME)
        print("模型加载完成。")
    except Exception as e:
        print(f"加载模型时发生错误: {e}")
        return

    # 3. 加载集合到内存以便搜索
    print(f"正在加载集合 '{COLLECTION_NAME}' 到内存...")
    try:
        collection.load()
        print("集合加载完成。")
    except Exception as e:
        print(f"加载集合时发生错误: {e}")
        return

    # 4. 循环接收用户输入并执行搜索
    print("\n---------- 开始问答检索 ----------")
    print("输入您的问题，或者输入 'quit' 或 'exit' 退出程序。")
    
    while True:
        try:
            query_text = input("\n请输入您的问题: ")
            if query_text.lower() in ["quit", "exit"]:
                print("程序退出。")
                break
            
            if not query_text.strip():
                print("输入不能为空，请重新输入。")
                continue

            # a. 将用户输入的问题文本转换为向量
            print("正在生成查询向量...")
            query_vector = model.encode(query_text)

            # b. 执行向量搜索
            print("正在检索相似问题...")
            results = collection.search(
                data=[query_vector],               # 要搜索的向量
                anns_field="question_embedding",   # 在哪个向量字段上搜索
                param=SEARCH_PARAMS,               # 搜索参数
                limit=TOP_K,                       # 返回最相似的 Top K 个结果
                output_fields=["question", "answer"] # 指定需要返回的字段
            )
            
            # c. 解析并显示结果
            print("\n---------- 检索结果 ----------")
            hit = results[0][0] # 获取第一个最相似的结果
            print(f"相似度 (距离): {hit.distance:.4f}")
            print(f"匹配到的问题: {hit.entity.get('question')}")
            print(f"对应的答案: {hit.entity.get('answer')}")
            print("------------------------------")

        except KeyboardInterrupt:
            print("\n检测到中断，程序退出。")
            break
        except Exception as e:
            print(f"在搜索过程中发生错误: {e}")

    # 5. 断开连接
    connections.disconnect("default")
    print("与 Milvus 的连接已关闭。")


if __name__ == '__main__':
    main() 